import math
from pathlib import Path
from typing import Dict

import cv2
import numpy as np

from detect import VehicleFasterRCNN
from optical_flow import SparseOpticalFlow
from track import SORT
from trajectory_structure import Trajectory


def draw_track_boxes(image, boxes):
    img_h, img_w, img_c = image.shape
    for i, box in enumerate(boxes):
        trk_id, x1, y1, x2, y2 = box
        color = (0, 0, 255)
        x1 = int(max(min(x1.item(), 1), 0) * img_w)
        y1 = int(max(min(y1.item(), 1), 0) * img_h)
        x2 = int(max(min(x2.item(), 1), 0) * img_w)
        y2 = int(max(min(y2.item(), 1), 0) * img_h)
        pt1, pt2 = (x1, y1), (x2, y2)
        cv2.rectangle(image, pt1, pt2, color, 2)
        cv2.putText(image, str(int(trk_id)), (x1 + 5, y1 + 15), 2, 0.5,
                    (255, 0, 0))
    return image


def draw_trajectories(image: np.ndarray, trajectories: Dict[int, Trajectory],
                      frame_index: int):
    color = (0, 0, 255)
    for tra_id, trajectory in trajectories.items():
        if len(trajectory.position_list) <= 0:
            continue
        if trajectory.last_frame_index < frame_index:
            continue
        tra_id = trajectory.tra_id
        box = trajectory.position_list[-1]
        x1, y1, x2, y2 = box
        pt1, pt2 = (x1, y1), (x2, y2)
        cv2.rectangle(image, pt1, pt2, color, 2)
        speed = 0
        if len(trajectory.speed_list) > 0:
            speed = math.floor(trajectory.speed_list[-1])
        if speed != 0:
            speed = F"{speed}km/h"
            cv2.putText(image, speed, (x1 + 5, y1 + 15), 1, 1, (0, 255, 255))
    return image


def filter_det_boxes(det_boxes: np.array) -> np.ndarray:
    """滤除长宽比不合适的检测框
    根据论文中的参数，beta_1 = 2.7, beta_2 = 1.2
    如果 w / h 或者 h / w 不在 [beta_2, beta_1] 范围内，则滤除此检测框
    """
    beta_1, beta_2 = 2.7, 1.2
    keep = list()
    for i, det_box in enumerate(det_boxes):
        x1, y1, x2, y2 = det_box
        w, h = x2 - x1, y2 - y1
        if beta_1 > w / h > beta_2 or beta_1 > h / w > beta_2:
            keep.append(i)
    det_boxes = det_boxes[keep]
    return det_boxes


def update_trajectories(trajectories: dict, trk_boxes: np.ndarray, H: np.array,
                        frame_index: int) -> None:
    for trk_box in trk_boxes:
        trk_id = int(trk_box[0].item())
        if trk_id not in trajectories:
            trajectories[trk_id] = Trajectory(tra_id=trk_id)
        else:
            trajectories[trk_id].add_position_and_moving_distance(
                trk_box, H, frame_index)


def detect_video():
    video_path = "data/videos/DJI_20200721145302_0001_WIDE.MP4"
    video = cv2.VideoCapture(video_path)
    if not video.isOpened():
        print(F"can not open the video: {video_path}.")
        return -1

    onnx_file_path = R"E:\models\ONNX\FasterRCNN-10.onnx"
    detector = VehicleFasterRCNN(onnx_file_path)
    tracker = SORT()
    optical = SparseOpticalFlow()

    # 以防视频开始帧不稳定，视频前几帧不考虑
    prve_frame = None
    prev_trk_boxes = None
    for _ in range(10):
        _, frame = video.read()
        det_boxes = detector.detect(frame)
        trk_boxes = tracker.track(det_boxes)
        prve_frame = frame.copy()
        prev_trk_boxes = trk_boxes.copy()

    car_trajectories = dict()

    window_name = "speed"
    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
    frame_index = 1

    video_save_path = Path(video_path).parent / \
        F"{Path(video_path).stem}_result.avi"
    video_save_path = str(video_save_path)
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    video_writer = cv2.VideoWriter(video_save_path, fourcc, 30.0, (1920, 1080))

    while True:
        _, frame = video.read()
        if frame is None:
            break

        det_boxes = detector.detect(frame)
        det_boxes = filter_det_boxes(det_boxes)

        trk_boxes = tracker.track(det_boxes)

        H = optical.optical_flow(prve_frame, frame, prev_trk_boxes)
        update_trajectories(car_trajectories, trk_boxes, H, frame_index)

        prve_frame = frame.copy()
        prev_trk_boxes = trk_boxes.copy()

        # draw_track_boxes(frame, trk_boxes)
        draw_trajectories(frame, car_trajectories, frame_index)
        cv2.imshow(window_name, frame)
        if cv2.waitKey(3) == 27:
            break
        frame_index += 1

        video_writer.write(frame)
    cv2.destroyWindow(window_name)

    # log_file = open("data/log.txt", "w")
    # for tra_id, trajectory in car_trajectories.items():
    #     log_file.writelines(F"{tra_id}: \n")
    #     for j, (diag_length, moving_d) in enumerate(
    #             zip(trajectory.diag_length_list,
    #                 trajectory.moving_distance_list)):
    #         log_file.writelines(F"{j+1}, {diag_length:.2f}, {moving_d:.2f} \n")
    #     log_file.writelines("-" * 20 + "\n")
    # log_file.close()


def main():
    detect_video()


if __name__ == "__main__":
    main()
